use crate::lexer::{Keyword,Token,SpannedToken,Symbol,Literal};
#[derive(Debug, Clone)]
pub enum AstNode {
    Program(Vec<AstNode>),
    VariableDecl { name: String, value: Option<Box<AstNode>> },
    FunctionDecl { name: String, params: Vec<String>, body: Vec<AstNode> },
    Expression(Expression),
    Statement(Statement),
    // 可扩展更多节点
}

#[derive(Debug, Clone)]
pub enum Expression {
    Identifier(String),
    Literal(String),
    BinaryOp { left: Box<Expression>, op: String, right: Box<Expression> },
    // 可扩展更多表达式
}

#[derive(Debug, Clone)]
pub enum Statement {
    Expr(Expression),
    If { cond: Expression, then_branch: Vec<AstNode>, else_branch: Option<Vec<AstNode>> },
    While { cond: Expression, body: Vec<AstNode> },
    Return(Option<Expression>),
    // 可扩展更多语句
}

pub struct Parser {
    tokens: Vec<SpannedToken>,
    pos: usize,
}

impl Parser {
    pub fn new(tokens: Vec<SpannedToken>) -> Self {
        Parser { tokens, pos: 0 }
    }

    fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.pos).map(|st| &st.token)
    }

    fn next(&mut self) -> Option<&Token> {
        let tok = self.tokens.get(self.pos).map(|st| &st.token);
        self.pos += 1;
        tok
    }

    pub fn parse(&mut self) -> Result<AstNode, String> {
        let mut nodes = Vec::new();
        while self.pos < self.tokens.len() {
            match self.parse_statement() {
                Ok(node) => nodes.push(node),
                Err(e) => return Err(e),
            }
        }
        Ok(AstNode::Program(nodes))
    }

    fn parse_statement(&mut self) -> Result<AstNode, String> {
        match self.peek() {
            Some(Token::Keyword(Keyword::Const)) => self.parse_variable_decl(),
            Some(Token::Keyword(Keyword::Fn)) => self.parse_function_decl(),
            Some(Token::Keyword(Keyword::If)) => self.parse_if_statement(),
            Some(Token::Keyword(Keyword::While)) => self.parse_while_statement(),
            Some(Token::Keyword(Keyword::Return)) => self.parse_return_statement(),
            _ => self.parse_expression_statement(),
        }
    }

    fn parse_variable_decl(&mut self) -> Result<AstNode, String> {
        self.next(); // consume 'const'
        let name = match self.next() {
            Some(Token::Identifier(n)) => n.clone(),
            _ => return Err(self.error_with_pos("Expected identifier after 'const'")),
        };
        let mut value = None;
        if let Some(Token::Symbol(Symbol::Equal)) = self.peek() {
            self.next(); // consume '='
            value = Some(Box::new(AstNode::Expression(self.parse_expression()?)));
        }
        // consume ';'
        if let Some(Token::Symbol(Symbol::Semicolon)) = self.peek() {
            self.next();
        }
        Ok(AstNode::VariableDecl { name, value })
    }

    fn parse_function_decl(&mut self) -> Result<AstNode, String> {
        self.next(); // consume 'fn'
        let name = match self.next() {
            Some(Token::Identifier(n)) => n.clone(),
            _ => return Err(self.error_with_pos("Expected function name")),
        };
        // consume '('
        if let Some(Token::Symbol(Symbol::LeftParen)) = self.peek() {
            self.next();
        } else {
            return Err(self.error_with_pos("Expected '(' after function name"));
        }
        let mut params = Vec::new();
        while let Some(Token::Identifier(param)) = self.peek() {
            params.push(param.clone());
            self.next();
            if let Some(Token::Symbol(Symbol::Comma)) = self.peek() {
                self.next();
            } else {
                break;
            }
        }
        // ')'
        if let Some(Token::Symbol(Symbol::RightParen)) = self.peek() {
            self.next();
        } else {
            return Err(self.error_with_pos("Expected ')' after parameters"));
        }
        // '{'
        if let Some(Token::Symbol(Symbol::LeftBrace)) = self.peek() {
            self.next();
        } else {
            return Err(self.error_with_pos("Expected '{' before function body"));
        }
        let mut body = Vec::new();
        while let Some(token) = self.peek() {
            if let Token::Symbol(Symbol::RightBrace) = token {
                break;
            }
            body.push(self.parse_statement()?);
        }
        // '}'
        if let Some(Token::Symbol(Symbol::RightBrace)) = self.peek() {
            self.next();
        }
        Ok(AstNode::FunctionDecl { name, params, body })
    }

    fn parse_expression_statement(&mut self) -> Result<AstNode, String> {
        let expr = self.parse_expression()?;
        // ';'
        if let Some(Token::Symbol(Symbol::Semicolon)) = self.peek() {
            self.next();
        }
        Ok(AstNode::Statement(Statement::Expr(expr)))
    }

    fn parse_expression(&mut self) -> Result<Expression, String> {
        match self.next() {
            Some(Token::Identifier(n)) => Ok(Expression::Identifier(n.clone())),
            Some(Token::Literal(lit)) => match lit {
                Literal::Integer(i) => Ok(Expression::Literal(i.to_string())),
                Literal::Float(f) => Ok(Expression::Literal(f.to_string())),
                Literal::String(s) => Ok(Expression::Literal(s.clone())),
                Literal::Char(c) => Ok(Expression::Literal(c.to_string())),
                Literal::Boolean(b) => Ok(Expression::Literal(b.to_string())),
                Literal::ByteString(bs) => Ok(Expression::Literal(format!("{:?}", bs))),
            },
            _ => Err(self.error_with_pos("Unexpected token in expression")),
        }
    }

    // 下面可以继续实现 if、while、return 等语句的解析
    fn parse_if_statement(&mut self) -> Result<AstNode, String> {
        // ...略
        Err(self.error_with_pos("Not implemented"))
    }
    fn parse_while_statement(&mut self) -> Result<AstNode, String> {
        // ...略
        Err(self.error_with_pos("Not implemented"))
    }
    fn parse_return_statement(&mut self) -> Result<AstNode, String> {
        // ...略
        Err(self.error_with_pos("Not implemented"))
    }
    fn error_with_pos(&self, msg: &str) -> String {
        if let Some(st) = self.tokens.get(self.pos) {
            format!("{} at line {}, column {}", msg, st.position.line, st.position.column)
        } else {
            format!("{} at end of file", msg)
        }
    }

}